
import os

import h5py as h5
import numpy as np


def load_dataset_label(file_path):
    with h5.File(file_path, "r") as f:
        dataset = np.array(f['dataset'])
        ground_truth = np.array(f['ground_truth'])

    return dataset, ground_truth


def load_datasets(file_name, test=False):
    with h5.File(file_name, "r") as f:
        dataMatrix = np.asarray(f["dataset"])

        if test:
            testDataset = np.asarray(f["test_dataset"])
            testKnn = np.asarray(f["test_knn"])

            result = (dataMatrix, testDataset, testKnn)

        else:
            trainDataset = np.asarray(f["train_dataset"])
            trainKnn = np.asarray(f["train_knn"])

            result = (dataMatrix, trainDataset, trainKnn)

        return result


def generate_dataset(file_name, target_dir='./dataset'):
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)

    if os.path.exists(file_name):
        dataset, ground_truth = load_dataset_label(file_name)
        nums = len(dataset)
        train_list = np.random.zipf(1.5, int(0.8 * len(dataset)))
        train_dataset = [dataset[index % nums] for index in train_list]
        train_knn = [ground_truth[index % nums] for index in train_list]

        test_list = np.random.zipf(1.5, int(0.2 * len(dataset)))
        test_dataset = [dataset[index % nums] for index in test_list]
        test_knn = [ground_truth[index % nums] for index in test_list]
        path, file_name = os.path.split(file_name)
        with h5.File(os.path.join(target_dir, file_name), "w") as f:
            f.create_dataset("dataset", data=dataset)
            f.create_dataset("train_dataset", data=train_dataset)
            f.create_dataset("train_knn", data=train_knn)
            f.create_dataset("test_dataset", data=test_dataset)
            f.create_dataset("test_knn", data=test_knn)

    else:
        print("数据集不存在：{}".format(file_name))


def calc_recall(t_index, p_index):
    """ 计算召回率

    :param t_index: 正确的结果
    :param p_index: 预测的结果
    :return: 召回率
    """
    t_index, p_index = list(t_index), list(p_index)

    t_index = set(t_index)
    p_index = set(p_index)

    tp = t_index & p_index
    tp_fn = t_index

    if tp_fn == 0:
        recall = 0
        print("没有正确结果")
    else:
        recall = len(tp) / len(tp_fn)

    return recall


if __name__ == '__main__':
    root_dir = 'E:'
    dir_path = os.path.join(root_dir, 'target_datasets')
    for file_name in os.listdir(dir_path):
        generate_dataset(os.path.join(dir_path, file_name))
